import numpy as np
import matplotlib.pyplot as plt
from matplotlib.font_manager import FontProperties
import os

# 设置中文字体
font_path = '/Users/liuyuzhuo/Library/Fonts/SimHei.ttf'  # 请确保此路径正确
font_prop = FontProperties(fname=font_path)
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

def plot_custom_confusion_matrix():
    """生成具有自定义数据的混淆矩阵"""
    # 创建自定义混淆矩阵数据
    cm = np.array([
        [200, 0],
        [0, 200]
    ])
    
    # 类别名称
    class_names = ['人脸', '花卉']
    
    # 创建图表
    fig, ax = plt.subplots(figsize=(10, 8))
    
    # 计算准确率
    total = np.sum(cm)
    accuracy = np.trace(cm) / total * 100
    
    # 创建标准化的混淆矩阵（按行归一化）
    cm_normalized = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
    
    # 设置颜色映射
    cmap = plt.cm.Blues
    
    # 绘制混淆矩阵
    im = ax.imshow(cm_normalized, interpolation='nearest', cmap=cmap)
    
    # 设置坐标轴标签
    tick_marks = np.arange(len(class_names))
    ax.set_xticks(tick_marks)
    ax.set_yticks(tick_marks)
    ax.set_xticklabels(class_names, fontproperties=font_prop, fontsize=14)
    ax.set_yticklabels(class_names, fontproperties=font_prop, fontsize=14)
    
    # 设置轴标签
    ax.set_ylabel('真实类别', fontproperties=font_prop, fontsize=16, labelpad=10)
    ax.set_xlabel('预测类别', fontproperties=font_prop, fontsize=16, labelpad=10)
    
    # 添加标题，包括准确率
    ax.set_title(f'混淆矩阵 (准确率: {accuracy:.2f}%)', fontproperties=font_prop, fontsize=18)
    
    # 添加文本标注到混淆矩阵单元格
    thresh = cm_normalized.max() / 2.
    for i in range(cm.shape[0]):
        for j in range(cm.shape[1]):
            text = "{}\n({:.1f}%)".format(cm[i, j], 100 * cm_normalized[i, j])
            ax.text(j, i, text,
                   ha="center", va="center",
                   color="white" if cm_normalized[i, j] > thresh else "black",
                   fontsize=14)
    
    # 添加颜色条
    fig.colorbar(im)
    
    # 确保输出目录存在
    os.makedirs("output/custom", exist_ok=True)
    
    # 保存图表
    plt.tight_layout()
    save_path = "output/custom/custom_confusion_matrix.png"
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    
    print(f"自定义混淆矩阵已保存至 {save_path}")

if __name__ == "__main__":
    plot_custom_confusion_matrix()
    print("完成！") 